use axum::{body::Body, http::Request, middleware::Next, response::Response};
use rust_mcp_transport::MCP_SESSION_ID_HEADER;
use std::{
    cell::RefCell,
    collections::HashMap,
    sync::{Mutex, OnceLock},
};
use tracing::info;

/// 域名信息结构体,存储在 Request extensions 中
#[derive(Clone, Debug)]
pub struct DomainInfo {
    pub domain: String,
}

/// 认证头信息结构体,存储在 Request extensions 中
#[derive(Clone, Debug)]
pub struct AuthHeaderInfo {
    pub header: String,
}

/// 提取Host header中间件
/// 从HTTP请求中提取Host header并存储到 Request extensions 中
pub async fn extract_host_middleware(
    mut request: Request<Body>,
    next: Next,
) -> Response {
    // Capture Authorization/Authentication headers for downstream use.
    let auth_header = request
        .headers()
        .get("authorization")
        .or_else(|| request.headers().get("authentication"))
        .and_then(|value| value.to_str().ok())
        .map(|value| value.to_string());

    let host = request
        .headers()
        .get("host")
        .and_then(|h| h.to_str().ok())
        .map(|h| h.split(':').next().unwrap_or(h).to_string());

    if let Some(domain) = host {
        info!("🎯 [MCP-MIDDLEWARE] 提取Host header域名: {}", domain);
        // 存储到 Request extensions 中
        request.extensions_mut().insert(DomainInfo {
            domain: domain.clone()
        });
        // 同步到任务/线程本地，供后续处理使用
        set_current_domain(domain.clone());

        if let Some(auth) = auth_header {
            request.extensions_mut().insert(AuthHeaderInfo {
                header: auth.clone()
            });
            set_current_auth_header(auth);
        }

        // 如果带有session-id，则把域名与session绑定，后续请求缺少Host时可兜底。
        if let Some(session_id) = request
            .headers()
            .get(MCP_SESSION_ID_HEADER)
            .and_then(|v| v.to_str().ok())
        {
            set_session_domain(session_id, &domain);
        }
    }

    let response = next.run(request).await;

    clear_current_domain();
    clear_current_auth_header();

    response
}

/// 从 Request extensions 中获取域名
pub fn get_domain_from_request<B>(request: &Request<B>) -> Option<String> {
    request
        .extensions()
        .get::<DomainInfo>()
        .map(|info| info.domain.clone())
}

/// 从 Request extensions 中获取认证头
pub fn get_auth_header_from_request<B>(request: &Request<B>) -> Option<String> {
    request
        .extensions()
        .get::<AuthHeaderInfo>()
        .map(|info| info.header.clone())
}

// 简单的任务/线程本地存储用于当前请求。
thread_local! {
    static CURRENT_DOMAIN: RefCell<Option<String>> = RefCell::new(None);
    static CURRENT_AUTH_HEADER: RefCell<Option<String>> = RefCell::new(None);
}

pub fn set_current_domain(domain: String) {
    CURRENT_DOMAIN.with(|d| {
        *d.borrow_mut() = Some(domain);
    });
}

pub fn get_current_domain() -> Option<String> {
    CURRENT_DOMAIN.with(|d| d.borrow().clone())
}

pub fn clear_current_domain() {
    CURRENT_DOMAIN.with(|d| {
        *d.borrow_mut() = None;
    });
}

pub fn set_current_auth_header(header: String) {
    CURRENT_AUTH_HEADER.with(|h| {
        *h.borrow_mut() = Some(header);
    });
}

pub fn get_current_auth_header() -> Option<String> {
    CURRENT_AUTH_HEADER.with(|h| h.borrow().clone())
}

pub fn clear_current_auth_header() {
    CURRENT_AUTH_HEADER.with(|h| {
        *h.borrow_mut() = None;
    });
}

// 全局 session -> domain 映射，跨线程兜底。
static SESSION_DOMAIN_MAP: OnceLock<Mutex<HashMap<String, String>>> = OnceLock::new();

fn session_domain_map() -> &'static Mutex<HashMap<String, String>> {
    SESSION_DOMAIN_MAP.get_or_init(|| Mutex::new(HashMap::new()))
}

/// 记录 session 与域名的映射，供跨线程请求兜底使用。
pub fn set_session_domain(session_id: &str, domain: &str) {
    let mut guard = session_domain_map().lock().unwrap();
    guard.insert(session_id.to_string(), domain.to_string());
}

/// 读取 session 绑定的域名。
pub fn get_domain_for_session(session_id: &str) -> Option<String> {
    let guard = session_domain_map().lock().unwrap();
    guard.get(session_id).cloned()
}

/// 清理 session 绑定的域名。
pub fn clear_session_domain(session_id: &str) {
    let mut guard = session_domain_map().lock().unwrap();
    guard.remove(session_id);
}